[graph_trainer] Add reusable dI/dW backward-node classification - #6

PR 1 of 2 (stack: this → graphpoc/deferred-dw)

Generalize the GraphPP dI/dW dependency analysis into a standalone helper the EP-overlap scheduler can consume: partition_backward_nodes classifies every node in the joint graph's gradient closures into di_nodes / dw_only_nodes / shared_nodes by pure ancestor algebra (closure of the input-gradient outputs vs closure of the parameter-gradient outputs), never inspecting op names or dtypes, and exposes movable_nodes, the subset of dW-only work a pass may DEFER (move later, never earlier).

Safety rules keep deferral provable:

• Pinned per node: collectives, mutations, CPU synchronizations (_local_scalar_dense or device-to-host reads), side-effectful ops, and producers with unclassified LIVE users; a pure dead subtree (e.g. the unused getitem outputs of a multi-output quantize op) does not pin its producer.
• Mutation hazards: a reader of a mutated buffer's alias family that sits BEFORE the write without a data edge through it is pinned (deferring it could push the read past the write); readers downstream of the write or after it keep their order against the fixed write. A write whose target cannot be resolved to graph values empties movable_nodes entirely.

split_di_dw.py stays untouched: its extraction-based dI membership deliberately includes dW-only placeholders as pass-through live-ins, and rebasing it on the closure helper would change which values it saves.

Validation

• Real-graph check on a tiny converted MXFP8 GroupedExperts joint trace (EP mesh wired, world size 1): all 3 wgrad GEMMs + 12 wgrad-chain quantize/rearrange nodes classify MOVABLE, dgrad chains stay on the dI side, A2A collectives are never movable; 210/230 nodes classified, movable == dw_only (48).
• 13 CPU unit tests (synthetic FX graphs, no GPU needed). black/usort/flake8 clean.
• Two real-graph-only bugs were caught by validating against the actual trace and are covered by the final design: dead getitems no longer pin their producer, and the global "mutation touching both closures" rule was replaced with per-reader hazard pinning (the dispatcher's index_put_-into-new_empty combine scatter is handled correctly).

🤖 Generated with Claude Code